#!/usr/bin/env python3
"""
添加用户区域关联表（多对多关系）
"""

import sqlite3
import sys
import os

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def migrate():
    """执行迁移"""
    db_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'instance', 'dispatch_system.db')
    
    if not os.path.exists(db_path):
        print(f"数据库文件不存在: {db_path}")
        return False
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        # 检查表是否已存在
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='user_areas'")
        if cursor.fetchone():
            print("user_areas表已存在，跳过创建")
            return True
        
        # 创建用户区域关联表
        cursor.execute("""
            CREATE TABLE user_areas (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                area_id INTEGER NOT NULL,
                FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE,
                FOREIGN KEY (area_id) REFERENCES areas (id) ON DELETE CASCADE,
                UNIQUE(user_id, area_id)
            )
        """)
        
        # 创建索引以提高查询性能
        cursor.execute("CREATE INDEX idx_user_areas_user_id ON user_areas(user_id)")
        cursor.execute("CREATE INDEX idx_user_areas_area_id ON user_areas(area_id)")
        
        # 迁移现有数据：将users表中的area_id关联迁移到user_areas表
        cursor.execute("""
            INSERT INTO user_areas (user_id, area_id)
            SELECT id, area_id FROM users 
            WHERE area_id IS NOT NULL
        """)
        
        conn.commit()
        print("user_areas表创建成功，现有数据已迁移")
        return True
        
    except Exception as e:
        conn.rollback()
        print(f"迁移失败: {e}")
        return False
    finally:
        conn.close()

if __name__ == '__main__':
    success = migrate()
    sys.exit(0 if success else 1)